from llmtcg.store.codebase.base import CodeBase
from llmtcg.store.codebase.c import CCodeBase
from loguru import logger

from argparse import ArgumentParser, Namespace
from time import time

import gc
import pickle


def parse_args() -> Namespace:
    parser = ArgumentParser("python3 -m llmtcg.store.codebase")
    subparsers = parser.add_subparsers(help="sub-command help")

    summary = subparsers.add_parser("summary")
    summary.add_argument("-l", "--language", type=str, choices={"cpp"}, default="cpp")
    summary.add_argument("-o", "--output", type=str, help="dump the summary file to a database")
    summary.add_argument("path", type=str, nargs="+")
    summary.set_defaults(func=summary_func)

    explore = subparsers.add_parser("explore")
    explore.add_argument("summary", type=str)
    explore.set_defaults(func=explore_func)

    return parser.parse_args()


def summary_func(args: Namespace) -> None:
    cb: CodeBase
    if args.language == 'cpp':
        cb = CCodeBase()
    else:
        raise NotImplementedError(args.language)

    for path in args.path:
        cb.load_directory(path)

    if args.output:
        with open(args.output, 'wb+') as f:
            pickle.dump(cb, f)


def explore_func(args: Namespace) -> None:
    start = time()
    with open(args.summary, "rb") as f:
        gc.disable()
        cb = pickle.load(f)
        gc.enable()

    cost = time() - start
    logger.info(f"summary loaded in {cost:.2f} seconds")
    breakpoint()


if __name__ == "__main__":
    args = parse_args()
    args.func(args)
